// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <memory>

#include <pollux/expression/eval_ctx.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/lib/check_duplicate_keys.h>
#include <pollux/functions/lib/rows_translation_util.h>
#include <pollux/vector/base_vector.h>
#include <pollux/vector/complex_vector.h>

namespace kumo::pollux::functions {
    namespace {
        static const char *kNullKeyErrorMessage = "map key cannot be null";
        static const char *kErrorMessageEntryNotNull = "map entry cannot be null";

        // See documentation at https://prestodb.io/docs/current/functions/map.html
        class MapFromEntriesFunction : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                POLLUX_CHECK_EQ(args.size(), 1);
                auto &arg = args[0];
                VectorPtr localResult;
                // Input can be constant or flat.
                if (arg->is_constant_encoding()) {
                    auto *constantArray = arg->as<ConstantVector<ComplexType> >();
                    const auto &flatArray = constantArray->value_vector();
                    const auto flatIndex = constantArray->index();

                    exec::LocalSelectivityVector singleRow(context, flatIndex + 1);
                    singleRow->clearAll();
                    singleRow->setValid(flatIndex, true);
                    singleRow->updateBounds();

                    localResult = applyFlat(
                        *singleRow.get(), flatArray->as<ArrayVector>(), outputType, context);
                    localResult =
                            BaseVector::wrap_in_constant(rows.size(), flatIndex, localResult);
                } else {
                    localResult =
                            applyFlat(rows, arg->as<ArrayVector>(), outputType, context);
                }

                context.moveOrCopyResult(localResult, rows, result);
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                return {
                    // unknown -> map(unknown, unknown)
                    exec::FunctionSignatureBuilder()
                    .returnType("map(unknown, unknown)")
                    .argumentType("unknown")
                    .build(),
                    // array(unknown) -> map(unknown, unknown)
                    exec::FunctionSignatureBuilder()
                    .returnType("map(unknown, unknown)")
                    .argumentType("array(unknown)")
                    .build(),
                    // array(row(K,V)) -> map(K,V)
                    exec::FunctionSignatureBuilder()
                    .typeVariable("K")
                    .typeVariable("V")
                    .returnType("map(K,V)")
                    .argumentType("array(row(K,V))")
                    .build()
                };
            }

        private:
            VectorPtr applyFlat(
                const SelectivityVector &rows,
                const ArrayVector *inputArray,
                const TypePtr &outputType,
                exec::EvalCtx &context) const {
                auto &inputValueVector = inputArray->elements();
                exec::LocalDecodedVector decodedRowVector(context);
                decodedRowVector.get()->decode(*inputValueVector);
                // If the input array(unknown) then all rows should have errors.
                if (inputValueVector->type_kind() == TypeKind::UNKNOWN) {
                    try {
                        POLLUX_USER_FAIL(kErrorMessageEntryNotNull);
                    } catch (...) {
                        context.setErrors(rows, std::current_exception());
                    }

                    auto sizes = allocateSizes(rows.end(), context.pool());
                    auto offsets = allocateSizes(rows.end(), context.pool());

                    // Output in this case is map(unknown, unknown), but all elements are
                    // nulls, all offsets and sizes are 0.
                    return std::make_shared<MapVector>(
                        context.pool(),
                        outputType,
                        inputArray->nulls(),
                        rows.end(),
                        sizes,
                        offsets,
                        BaseVector::create(UNKNOWN(), 0, context.pool()),
                        BaseVector::create(UNKNOWN(), 0, context.pool()));
                }

                exec::LocalSelectivityVector remianingRows(context, rows);
                auto row_vector = decodedRowVector->base()->as<RowVector>();
                auto keyVector = row_vector->childAt(0);

                BufferPtr sizes = allocateSizes(rows.end(), context.pool());
                vector_size_t *mutableSizes = sizes->asMutable<vector_size_t>();
                rows.applyToSelected([&](vector_size_t row) {
                    mutableSizes[row] = inputArray->rawSizes()[row];
                });

                auto resetSize = [&](vector_size_t row) { mutableSizes[row] = 0; };

                // Validate all map entries and map keys are not null.
                if (decodedRowVector->may_have_nulls() || keyVector->may_have_nulls() ||
                    keyVector->may_have_nulls_recursive()) {
                    context.applyToSelectedNoThrow(rows, [&](vector_size_t row) {
                        const auto size = inputArray->sizeAt(row);
                        const auto offset = inputArray->offsetAt(row);

                        for (auto i = 0; i < size; ++i) {
                            // Check nulls in the top level row vector.
                            const bool isMapEntryNull = decodedRowVector->is_null_at(offset + i);
                            if (isMapEntryNull) {
                                // Set the sizes to 0 so that the final map vector generated is
                                // valid in case we are inside a try. The map vector needs to be
                                // valid because its consumed by checkDuplicateKeys before try
                                // sets invalid rows to null.
                                resetSize(row);
                                POLLUX_USER_FAIL(kErrorMessageEntryNotNull);
                            }

                            // Check null keys.
                            auto keyIndex = decodedRowVector->index(offset + i);
                            if (keyVector->is_null_at(keyIndex)) {
                                resetSize(row);
                                POLLUX_USER_FAIL(kNullKeyErrorMessage);
                            }
                        }
                    });
                }

                context.deselectErrors(*remianingRows.get());

                VectorPtr wrappedKeys;
                VectorPtr wrappedValues;
                if (decodedRowVector->isIdentityMapping()) {
                    wrappedKeys = row_vector->childAt(0);
                    wrappedValues = row_vector->childAt(1);
                } else if (decodedRowVector->isConstantMapping()) {
                    if (decodedRowVector->is_null_at(0)) {
                        // If top level row is null, child might not be addressable at index 0
                        // so we do not try to read it.
                        wrappedKeys = BaseVector::create_null_constant(
                            row_vector->childAt(0)->type(),
                            decodedRowVector->size(),
                            context.pool());
                        wrappedValues = BaseVector::create_null_constant(
                            row_vector->childAt(1)->type(),
                            decodedRowVector->size(),
                            context.pool());
                    } else {
                        wrappedKeys = BaseVector::wrap_in_constant(
                            decodedRowVector->size(),
                            decodedRowVector->index(0),
                            row_vector->childAt(0));
                        wrappedValues = BaseVector::wrap_in_constant(
                            decodedRowVector->size(),
                            decodedRowVector->index(0),
                            row_vector->childAt(1));
                    }
                } else {
                    // Dictionary.
                    auto indices = allocate_indices(decodedRowVector->size(), context.pool());
                    auto nulls = allocate_nulls(decodedRowVector->size(), context.pool());
                    auto *mutable_nulls = nulls->asMutable<uint64_t>();
                    memcpy(
                        indices->asMutable<vector_size_t>(),
                        decodedRowVector->indices(),
                        BaseVector::byteSize<vector_size_t>(decodedRowVector->size()));
                    // Any null in the top row(X, Y) should be marked as null since its
                    // not guranteed to be addressable at X or Y.
                    for (auto i = 0; i < decodedRowVector->size(); i++) {
                        if (decodedRowVector->is_null_at(i)) {
                            bits::set_null(mutable_nulls, i);
                        }
                    }
                    wrappedKeys = BaseVector::wrap_in_dictionary(
                        nulls, indices, decodedRowVector->size(), row_vector->childAt(0));
                    wrappedValues = BaseVector::wrap_in_dictionary(
                        nulls, indices, decodedRowVector->size(), row_vector->childAt(1));
                }

                // To avoid creating new buffers, we try to reuse the input's buffers
                // as many as possible.
                auto map_vector = std::make_shared<MapVector>(
                    context.pool(),
                    outputType,
                    inputArray->nulls(),
                    rows.end(),
                    inputArray->offsets(),
                    sizes,
                    wrappedKeys,
                    wrappedValues);

                checkDuplicateKeys(map_vector, *remianingRows, context);
                return map_vector;
            }
        };
    } // namespace

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_map_from_entries,
        MapFromEntriesFunction::signatures(),
        std::make_unique<MapFromEntriesFunction>());
} // namespace kumo::pollux::functions
